#!/usr/bin/python3
# -*- coding: utf-8 -*-

# @file engine.py
# @brief
# @author QRS
# @version 1.0
# @date 2022-02-18 11:56

import numpy as np
from repnet import (
    get_repnet_model,
    get_counts
)

from frepai.core.base import TaskBase
from frepai.core.message import ServiceType
from frepai.constants import ENGINE_TOPIC


class EngineTask(TaskBase):
    def __init__(self, logger, handler):
        super(EngineTask, self).__init__(logger, handler, ServiceType.S_ENGINE, ENGINE_TOPIC)
        self.ckpts = '/ckpts'

    def init(self):
        self.model = get_repnet_model(self.ckpts)

    def step_run(self, cfg, H):
        self.logger.info(f'engine: {cfg}')
        workdir = cfg['workdir']
        frames = np.load(f'{workdir}/model_inputs.npy')

        def _progress(x):
            self.logger.info(x)

        (pred_period, pred_score,
            within_period, per_frame_counts,
            chosen_stride, final_embs, feature_maps, feat_factors) = get_counts(self.model, frames,
                cfg.get('strides', [1]),
                cfg.get('batch_size', 20),
                progress_cb=_progress)
